package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
)

// AesEncrypt 加密
func AesEncrypt(keyStr string, data []byte) (string, error) {
	// 把秘钥转成数组
	key := []byte(keyStr)
	// 创建加密实例
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	blockSize := block.BlockSize()
	// 通过pkcs7Padding(data, blockSize)对数据进行填充，确保数据长度为块大小的整数倍
	// 填充规则：缺n字节则填充n个0x0n，例如缺5字节填充0x05重复5次
	encryptBytes := pkcs7Padding(data, blockSize)
	result := make([]byte, len(encryptBytes))
	// 使用 cbc 加密模式
	// 使用cipher.NewCBCEncrypter创建加密器，这里将密钥前blockSize字节作为初始化向量(IV)
	blockMode := cipher.NewCBCEncrypter(block, key[:blockSize])
	// blockMode.CryptBlocks完成分组加密，结果存入result字节数组
	// 最终通过Base64编码输出字符串
	blockMode.CryptBlocks(result, encryptBytes)
	return base64.StdEncoding.EncodeToString(result), nil
}

// AesDecrypt 解密
func AesDecrypt(keyStr string, dataStr string) ([]byte, error) {
	// 创建实例
	key := []byte(keyStr)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	// 先对输入字符串进行Base64解码base64.StdEncoding.DecodeString
	data, err := base64.StdEncoding.DecodeString(dataStr)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	// 使用与加密相同的密钥创建AES实例
	// 通过cipher.NewCBCDecrypter创建CBC模式解密器，同样复用密钥前段作为IV
	blockMode := cipher.NewCBCDecrypter(block, key[:blockSize])
	result := make([]byte, len(data))
	// blockMode.CryptBlocks完成分组解密，结果存入result字节数组
	blockMode.CryptBlocks(result, data)
	// 调用pkcs7UnPadding移除填充字节，通过最后一个字节的值确定需截断的字节数
	result, err = pkcs7UnPadding(result)
	if err != nil {
		return nil, err
	}
	return result, nil
}

// 该函数用于在加密前对数据进行填充，使其长度符合分组加密算法的块大小要求。
func pkcs7Padding(data []byte, blockSize int) []byte {
	// 通过 blockSize - len(data) % blockSize 确定需要填充的字节数。
	// 例如，若数据长度差3字节不足块大小16字节，则填充3个0x03；若数据恰好对齐块大小（如16字节），则填充16个0x10
	padding := blockSize - len(data)%blockSize
	// 使用 bytes.Repeat 生成由重复填充值组成的字节序列，例如缺5字节时填充5个0x05
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// 该函数用于解密后移除填充数据，恢复原始内容。
func pkcs7UnPadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, errors.New("empty encrypt data")
	}
	// 通过读取数据最后一个字节的值（unPadding），确定填充的字节数。例如，若最后一个字节是0x05，则需移除末尾5字节
	unPadding := int(data[length-1])
	// 通过 data[:(length - unPadding)] 截断填充部分，返回原始数据。
	return data[:(length - unPadding)], nil
}

// Sha256 加密
func Sha256(data string) string {
	hash := sha256.New()
	_, err := io.WriteString(hash, data)
	if err != nil {
		return ""
	}
	hashValue := hash.Sum(nil)
	return fmt.Sprintf("%x", hashValue)
}

// MD5 加密
func MD5(data string) string {
	md5bs := md5.Sum([]byte(data))
	return hex.EncodeToString(md5bs[:])
}
